from dataclasses import dataclass
from typing import Dict, List, Any, Tuple, Optional
import queue


from high_rotary_evaporation_os.GlobalVariables import *

# ======= Device State Machines =======
class DeviceStateMachine:
    def __init__(self, name: str, global_vars: GlobalVariables):
        self.name = name
        self.global_vars = global_vars
        self.logger = logging.getLogger(f"{name}_state_machine")

    def update_status(self, status: int) -> None:
        """Update the status of this device"""
        self.global_vars.update_status(f"{self.name}_status", status)

    def get_status(self) -> int:
        """Get the current status of this device"""
        return self.global_vars.get_status(f"{self.name}_status")


class ChromatographModule(DeviceStateMachine):
    def __init__(self, global_vars: GlobalVariables):
        super().__init__("chromatograph_module", global_vars)

    def start_collection(self) -> None:
        """Start collecting samples"""
        self.update_status(ChromatographModuleStatus.COLLECTING)

    def complete_collection(self) -> None:
        """Complete collection"""
        self.update_status(ChromatographModuleStatus.WAITING_MERGE)

    def start_merging(self) -> None:
        """Start merging samples"""
        self.update_status(ChromatographModuleStatus.MERGING)

    def complete_merging(self) -> None:
        """Complete merging"""
        self.update_status(ChromatographModuleStatus.COMPLETED)

    def start_separation(self) -> None:
        """Start separation process"""
        self.update_status(ChromatographModuleStatus.SEPARATING)

    def complete_separation(self) -> None:
        """Complete separation"""
        self.update_status(ChromatographModuleStatus.IDLE)


class RotaryEvaporation(DeviceStateMachine):
    def __init__(self, global_vars: GlobalVariables):
        super().__init__("rotary_evaporation", global_vars)

    def start_evaporation(self) -> None:
        """Start evaporation process"""
        self.update_status(RotaryEvaporationStatus.EVAPORATING)

    def complete_evaporation(self) -> None:
        """Complete evaporation"""
        self.update_status(RotaryEvaporationStatus.WAITING_DISCHARGE)

    def discharge_waste(self) -> None:
        """Discharge waste"""
        self.update_status(RotaryEvaporationStatus.IDLE)

    def wait_evaporation(self) -> None:
        """Wait for evaporation"""
        self.update_status(RotaryEvaporationStatus.WAITING_EVAPORATION)


# Similar classes for other devices...

# ======= Task Scheduler =======
@dataclass
class Task:
    name: str
    precondition: Dict[str, int]  # Component name and required status
    action: str
    run_status: Dict[str, int]  # Status during running
    end_status: Dict[str, int]  # Status after completion
    priority: int = 999   # 新增，默认999，数字越小优先级越高
    description: str = ""
    notes: str = ""


class TaskScheduler:
    def __init__(self, global_vars: GlobalVariables,action_executor ):
        self.global_vars = global_vars
        self.tasks = []
        self.task_queue = queue.Queue()
        self.is_running = False
        self.logger = logging.getLogger("task_scheduler")
        self.action_executor = action_executor  # 新增


    def add_task(self, task: Task) -> None:
        """Add a task to the scheduler"""
        self.tasks.append(task)
        self.logger.info(f"Added task: {task.name}")

    def initialize_tasks(self) -> None:
        """Initialize tasks from the table definitions"""
        # Add tasks from the first image (task definition table)



        # Task 1: 加样 (sample injection)
        self.add_task(Task(
            name="sample_injection",
            precondition={
                "column_rack_status": ColumnRackStatus.IN_USE,
                "ethyl_acetate_status": EthylAcetateStatus.SUFFICIENT,
                "petroleum_ether_status": PetroleumEtherStatus.SUFFICIENT,
                "scara_status": ScaraStatus.IDLE
            },
            action="inject_sample",
            run_status={"scara_status": ScaraStatus.IN_USE},
            end_status={"scara_status": ScaraStatus.IDLE},
            priority = 10

        ))


        # Task 2: 收集 (collection)
        self.add_task(Task(
            name="collection",
            precondition={
                "chromatograph_module_status": ChromatographModuleStatus.IDLE,
                "column_rack_status": ColumnRackStatus.IN_USE,
                "ethyl_acetate_status": EthylAcetateStatus.SUFFICIENT,
                "petroleum_ether_status": PetroleumEtherStatus.SUFFICIENT
            },
            action="collect_sample",
            run_status={"chromatograph_module_status": ChromatographModuleStatus.COLLECTING},
            end_status={"chromatograph_module_status": ChromatographModuleStatus.WAITING_MERGE},
            priority = 2

        ))

        # Task 3: 混合 (mixing)
        self.add_task(Task(
            name="mixing",
            precondition={
                "chromatograph_module_status": ChromatographModuleStatus.WAITING_MERGE,
                "ethyl_acetate_status": EthylAcetateStatus.SUFFICIENT,
                "petroleum_ether_status": PetroleumEtherStatus.SUFFICIENT
            },
            action="mix_sample",
            run_status={"chromatograph_module_status": ChromatographModuleStatus.MERGING},
            end_status={"chromatograph_module_status": ChromatographModuleStatus.IDLE},
            priority=1
        ))

        # Task 3: 混合 (mixing)
        self.add_task(Task(
            name="mixing",
            precondition={
                "chromatograph_module_status": ChromatographModuleStatus.WAITING_MERGE,
                "ethyl_acetate_status": EthylAcetateStatus.SUFFICIENT,
                "petroleum_ether_status": PetroleumEtherStatus.SUFFICIENT
            },
            action="mix_sample",
            run_status={"chromatograph_module_status": ChromatographModuleStatus.MERGING},
            end_status={"chromatograph_module_status": ChromatographModuleStatus.IDLE},
            priority=1
        ))



    def check_preconditions(self, task: Task) -> bool:
        """Check if a task's preconditions are met"""
        all_met = True
        for component, required_status in task.precondition.items():
            current_status = self.global_vars.get_status(component)
            if current_status != required_status:
                self.logger.info(
                    f"[{task.name}] 前置条件未满足：{component} 当前状态={current_status}，期望={required_status}"
                )
                all_met = False
            else:
                self.logger.info(
                    f"[{task.name}] 前置条件已满足：{component} 当前状态={current_status}"
                )
        return all_met

    def update_run_status(self, task: Task) -> None:
        """Update the status of components during task execution"""
        for component, status in task.run_status.items():
            self.global_vars.update_status(component, status)
            self.logger.info(f"[{task.name}] 设置运行中状态：{component} => {status}")

    def update_end_status(self, task: Task) -> None:
        """Update the status of components after task completion"""
        for component, status in task.end_status.items():
            self.global_vars.update_status(component, status)
            self.logger.info(f"[{task.name}] 设置结束后状态：{component} => {status}")

    def queue_eligible_tasks(self) -> None:
        """Queue all tasks whose preconditions are met"""
        # 按优先级升序排列（数字越小越高优先级）
        sorted_tasks = sorted(self.tasks, key=lambda t: t.priority)
        for task in sorted_tasks:
            if self.check_preconditions(task):
                self.task_queue.put(task)
                self.logger.info(f"Queued task: {task.name} (priority={task.priority})")
                return  # 只入队第一个找到的满足条件的任务
        self.logger.info("当前无可执行任务。")

    def execute_task(self, task: Task) -> None:
        """Execute a task"""
        self.logger.info(f"Executing task: {task.name}")
        self.update_run_status(task)

        # Simulate task execution
        time.sleep(2)  # Simulate task duration

        # === 动态调用动作 ===
        if self.action_executor and hasattr(self.action_executor, task.action):
            self.logger.info(f"Calling action_executor method: {task.action}")
            getattr(self.action_executor, task.action)()
        else:
            self.logger.warning(f"No action_executor method for {task.action}")

            # 阻塞等待人工干预
            print(f"\n[警告] 未找到 {task.action} 的动作实现，系统已暂停。")
            print("请处理异常并确认后输入 'c' 继续，或 'q' 退出系统。")
            while True:
                user_input = input("输入 c 继续，q 退出：").strip().lower()
                if user_input == 'c':
                    self.logger.info("用户确认继续执行后续任务。")
                    break
                elif user_input == 'q':
                    self.logger.info("用户选择退出，停止任务调度。")
                    self.is_running = False
                    return
                else:
                    print("输入无效，请输入 'c' 或 'q'。")

        time.sleep(2)  # 模拟任务耗时

        self.update_end_status(task)
        self.logger.info(f"Completed task: {task.name}")

    def run(self) -> None:
        """Run the task scheduler"""
        self.is_running = True
        while self.is_running:
            # Queue eligible tasks
            self.queue_eligible_tasks()

            # Execute tasks from the queue
            try:
                task = self.task_queue.get(block=False)
                self.execute_task(task)
            except queue.Empty:
                # No tasks to execute
                time.sleep(1)